use anyhow::{ anyhow, Result};
use std::{
    thread,
    fmt, 
    sync::mpsc,
    ops::{Add, AddAssign, Mul},
};
use crate::{dot_product, Vector};
// 矩阵的乘法
// [[],[],[]]

const THREAD_NUM: usize = 4; 

//#[derive(Debug)]
pub struct Matrix<T: fmt::Debug> {
   pub data: Vec<T>,
   pub row: usize,
   pub col: usize,

}

pub struct MsgInput<T>{
    idx: usize,
    row: Vector<T>,
    col: Vector<T>,
}

pub struct MsgOutput<T>{
    idx: usize,
    value: T, 
}

pub struct Msg<T> {
   pub input: MsgInput<T>,
   pub sender: oneshot::Sender<MsgOutput<T>>,
}

pub fn multiply<T>(a: &Matrix<T>, b: &Matrix<T>) -> Result<Matrix<T>>
where T: fmt::Debug + Add<Output = T> + AddAssign + Mul<Output = T>  + Copy +Default + Send  + 'static,  // T不能安全地在多线程中传递， 就用到了这里的 send + 'static
{
    if a.col != b.row {
        return Err(anyhow!(" matrix  multiply error  a.col != b.row "));
    }

    // generate Thread_num  thread to do dot_product
    let senders =  (0..THREAD_NUM)
    .map(|_| {
        let (tx,rx) = mpsc::channel::<Msg<T>>();
        thread::spawn(move || {
            for msg in rx {
                let value = dot_product(msg.input.row, msg.input.col)?;
                if let Err(e) = msg.sender.send(MsgOutput {
                    idx: msg.input.idx,
                    value,
                }){
                  eprintln!("send error: {:?}",e) ;
                };
            }
            Ok::<_, anyhow::Error>(())
        });
        tx

    })
    .collect::<Vec<_>>();

    let mut data = vec![T::default(); a.row  * b.col];//Vec::with_capacity(a.row  * b.col);
    let mut receivers = Vec::with_capacity(a.row * b.col);


    // map- reduce =>  map phase
    for i in 0..a.row{
        for j in 0..b.col{
            let row = Vector::new(&a.data[i*a.col..(i+1) * a.col]);
            let col_data = b.data[j..]
            .iter()
            .step_by(b.col)
            .copied()
            .collect::<Vec<_>>();
            let col = Vector::new(col_data);
            
            let idx = i*b.col +j ;
            let input = MsgInput::new(  idx, row, col);
            let (tx,rx) = oneshot::channel();

            let msg = Msg::new(input,tx);
            if let Err(e) = senders[idx % THREAD_NUM ].send(msg){
                eprintln!("sende error: {:?}",e);
            }

            receivers.push(rx);
            // data[i*b.col +j] += dot_product(row, col)?;
          


            // for k in 0..a.col {
            //     data[i*b.col +j ] += a.data[i*a.col +k] * b.data[k*b.col +j];
            // }   
        }
    }

    // map- reduce =>  reduce phase
    for rx in receivers{
        let output = rx.recv()?;
        data[output.idx] = output.value;
    }

    Ok(Matrix { data, row: a.row, col: b.col })

}


// 主是要可以以Debug 方式 打印出来
impl<T: fmt::Debug> Matrix<T> {
    pub fn new(data: impl Into<Vec<T>>, row: usize,col: usize) -> Self {
        Self { data: data.into(), row, col }
    }
}

impl<T> Mul for Matrix<T>
where T: fmt::Debug + Add<Output = T> + AddAssign + Mul<Output = T>  + Copy +Default + Send  + 'static, 
{
    type Output = Self;

    fn mul(self, rhs: Self) -> Self::Output {
        multiply(&self, &rhs).expect("matrix multiply error")
    }
}

//不用Debug  ， 也可以实现Display  trait
impl<T> fmt::Display for Matrix<T> 
where T: fmt::Display  + fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for i in 0..self.row{
            write!(f,"{{")?;
            for j in 0..self.col {
                write!(f,"{:?}",self.data[i*self.col +j])?;
                if j != self.col-1{
                    write!(f, " ")?;
                }
            }
            
            write!(f,"}}")?;
        }

        Ok(())
    }
}


impl<T> fmt::Debug for Matrix<T>
where T: fmt::Display + fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
       // f.debug_struct("Matrix").field("data", &self.data).field("row", &self.row).field("col", &self.col).finish()
       write!(f,
    "Matrix {{data : {:?}, row: {}, col: {} }}",
      self.data, self.row,self.col
      )
    }
}

impl<T> MsgInput<T> {
    pub fn new(idx: usize, row: Vector<T> , col: Vector<T> ) ->Self {
        Self{idx,row,col}
    }
}

impl<T> Msg<T> {
    pub fn new(input : MsgInput<T>, sender: oneshot::Sender<MsgOutput<T>>) -> Self {
        Self{input, sender}
    }
}




#[cfg(test)]
mod tests{
    use anyhow::Ok;

    use super::*;

    #[test]

    fn test_matrix_multiply() ->Result<()> {
        // let a = Matrix::new([1,2,3,4,5,6], 2,3);
        // let b = Matrix::new([1,2,3,4,5,6], 3,2);
        let a = Matrix::new([1,2,3,4], 2,2);
        let b = Matrix::new([1,2,3,4], 2,2);
        // let c = multiply(&a, &b);
        let c = a * b;
        //format!("{:#?}",c);
       assert_eq!(c.data, vec![7,9,11,13]);
       assert_eq!(format!("{}",c), "{7 9}{11 13}");
     //   assert_eq!(c.unwrap().col,2);
     //assert_eq!(c.col,2);
        Ok(())
    }


    #[test]
    #[should_panic]
    
    fn test_should_panic() {
        let a = Matrix::new([1,2,3,4,5,6], 2,3);
        let b = Matrix::new([1,2,3,4], 2,2);
        a * b;
    }
    

    #[test]
    fn test_a_can_not_mutiply_b() {
        let a = Matrix::new([1,2,3,4,5,6], 2,3);
        let b = Matrix::new([1,2,3,4], 2,2);
        let c = multiply(&a, &b);
        assert!(c.is_err());
    }


}
